package org.lds.hotkey.common.autoconfigure.etcd;

import io.etcd.jetcd.ByteSequence;
import io.etcd.jetcd.KV;
import io.etcd.jetcd.KeyValue;
import io.etcd.jetcd.Watch;
import io.etcd.jetcd.kv.DeleteResponse;
import io.etcd.jetcd.kv.GetResponse;
import io.etcd.jetcd.options.GetOption;
import io.etcd.jetcd.options.PutOption;
import io.etcd.jetcd.options.WatchOption;
import io.etcd.jetcd.watch.WatchResponse;
import lombok.extern.slf4j.Slf4j;

import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.function.Consumer;
import java.util.stream.Collectors;

@Slf4j
public class EtcdService {

    private final EtcdClient etcdClient;

    public EtcdService(EtcdClient etcdClient) {
        this.etcdClient = etcdClient;
    }

    private ByteSequence fromDefault(String str) {
        return ByteSequence.from(str, StandardCharsets.UTF_8);
    }

    public String get(String key) {
        try {
            ByteSequence byteKey = fromDefault(key);
            GetResponse response = etcdClient.getKv().get(byteKey).get();
            List<KeyValue> values = response.getKvs();
            if (values.isEmpty()) return null;
            return values.get(0).getValue().toString();
        } catch (Throwable e) {
            log.error("get fail.", e);
            return null;
        }
    }

    public List<String> listValueByPrefix(String key) {
        try {
            ByteSequence byteKey = fromDefault(key);
            GetOption option = GetOption.newBuilder().isPrefix(true).build();
            GetResponse response = etcdClient.getKv().get(byteKey, option).get();
            List<KeyValue> values = response.getKvs();
            if (values.isEmpty()) return null;
            return values.stream().map(keyValue -> keyValue.getValue().toString()).collect(Collectors.toList());
        } catch (Throwable e) {
            log.error("getByPrefix fail.", e);
            return null;
        }
    }

    public List<KeyValue> listByPrefix(String key) {
        try {
            ByteSequence byteKey = fromDefault(key);
            GetOption option = GetOption.newBuilder().isPrefix(true).build();
            GetResponse response = etcdClient.getKv().get(byteKey, option).get();
            return response.getKvs();
        } catch (Throwable e) {
            log.error("getByPrefix fail.", e);
            return null;
        }
    }

    public boolean put(String key, String value) {
        try {
            ByteSequence byteKey = fromDefault(key);
            ByteSequence byteValue = fromDefault(value);
            KV kv = etcdClient.getKv();
            kv.put(byteKey, byteValue).get();
            return true;
        } catch (Throwable e) {
            log.error("put fail.", e);
            return false;
        }
    }

    public boolean put(String key, String value, long timeout) {
        try {
            long leaseId = etcdClient.getLease().grant(timeout).get().getID();
            PutOption option = PutOption.newBuilder().withLeaseId(leaseId).build();

            ByteSequence byteKey = fromDefault(key);
            ByteSequence byteValue = fromDefault(value);
            etcdClient.getKv().put(byteKey, byteValue, option).get();
            return true;
        } catch (Throwable e) {
            log.error("put fail.", e);
            return false;
        }
    }

    public boolean delete(String key) {
        try {
            ByteSequence byteKey = fromDefault(key);
            etcdClient.getKv().delete(byteKey).get();
            return true;
        } catch (Exception e) {
            log.error("delete fail.", e);
            return false;
        }
    }

    public boolean deleteAndRevoke(String key) {
        try {
            ByteSequence byteKey = fromDefault(key);
            DeleteResponse deleteResponse = etcdClient.getKv().delete(byteKey).get();
            long leaseId = deleteResponse.getPrevKvs().get(0).getLease();
            etcdClient.getLease().revoke(leaseId);
            return true;
        } catch (Throwable e) {
            log.error("deleteAndRevoke fail.", e);
            return false;
        }
    }

    /**
     * 监听key
     *
     * @param key
     * @param consumer
     * @return
     */
    public Watch.Watcher watchKey(String key, Consumer<WatchResponse> consumer) {
        try {
            ByteSequence watchKey = fromDefault(key);
            WatchOption option = WatchOption.newBuilder().withPrefix(watchKey).withPrevKV(true).build();
            return etcdClient.getWatch().watch(watchKey, option, consumer);
        } catch (Throwable e) {
            log.error("watchKey fail.", e);
            return null;
        }
    }

}
